package core

import (
	"errors"
	"sync"

	"github.com/cogentcore/webgpu/wgpu"
)

// MemoryPool 实现GPU内存池，按usage分类
type MemoryPool struct {
	device   *wgpu.Device
	buckets  map[uint64]map[wgpu.BufferUsage][]*wgpu.Buffer // 按size和usage双重索引
	mutex    sync.RWMutex
	maxSize  uint64
	usedSize uint64
	disabled bool // ✅ 新增：标记是否禁用池化
}

// NewMemoryPool 创建内存池
func NewMemoryPool(device *wgpu.Device, maxSize uint64) *MemoryPool {
	return &MemoryPool{
		device:   device,
		buckets:  make(map[uint64]map[wgpu.BufferUsage][]*wgpu.Buffer),
		maxSize:  maxSize,
		disabled: maxSize == 0, // ✅ 当maxSize为0时禁用池化
	}
}

// getAlignedSize 对齐到2的幂次方
func (p *MemoryPool) getAlignedSize(size uint64) uint64 {
	size--
	size |= size >> 1
	size |= size >> 2
	size |= size >> 4
	size |= size >> 8
	size |= size >> 16
	size++
	return size
}

// Get 获取buffer（优先复用），现在会匹配usage
func (p *MemoryPool) Get(size uint64, usage wgpu.BufferUsage) (*wgpu.Buffer, error) {
	// ✅ 禁用模式：直接创建新buffer
	if p.disabled {
		desc := &wgpu.BufferDescriptor{
			Size:  size,
			Usage: usage,
			Label: "direct-buffer",
		}
		return p.device.CreateBuffer(desc)
	}

	p.mutex.Lock()
	defer p.mutex.Unlock()

	alignedSize := p.getAlignedSize(size)

	// 尝试复用相同size和usage的buffer
	if usageBuckets, ok := p.buckets[alignedSize]; ok {
		if buffers, ok := usageBuckets[usage]; ok && len(buffers) > 0 {
			buf := buffers[len(buffers)-1]
			usageBuckets[usage] = buffers[:len(buffers)-1]
			p.usedSize += alignedSize
			return buf, nil
		}
	}

	// 检查内存限制
	if p.usedSize+alignedSize > p.maxSize {
		return nil, errors.New("memory pool exceeded")
	}

	// 创建新buffer
	desc := &wgpu.BufferDescriptor{
		Size:  alignedSize,
		Usage: usage,
		Label: "pool-buffer",
	}

	buffer, err := p.device.CreateBuffer(desc)
	if err != nil {
		return nil, err
	}

	p.usedSize += alignedSize
	return buffer, nil
}

// 完整方案：修改Put方法签名
func (p *MemoryPool) Put(buffer *wgpu.Buffer, size uint64, usage wgpu.BufferUsage) {
	// ✅ 禁用模式：直接释放
	if p.disabled {
		buffer.Release()
		return
	}

	p.mutex.Lock()
	defer p.mutex.Unlock()

	alignedSize := p.getAlignedSize(size)
	p.usedSize -= alignedSize

	// 按size和usage分类存储
	if p.buckets[alignedSize] == nil {
		p.buckets[alignedSize] = make(map[wgpu.BufferUsage][]*wgpu.Buffer)
	}
	p.buckets[alignedSize][usage] = append(p.buckets[alignedSize][usage], buffer)
}

// Close 释放所有池化内存
func (p *MemoryPool) Close() error {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	for _, usageBuckets := range p.buckets {
		for _, buffers := range usageBuckets {
			for _, buf := range buffers {
				buf.Release()
			}
		}
	}
	p.buckets = nil
	p.usedSize = 0
	return nil
}

// Stats 返回内存池统计
func (p *MemoryPool) Stats() (total, used uint64) {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	return p.maxSize, p.usedSize
}
